from cam import Camera
from flask import Flask, Response
import cv2
import numpy as np
import infer_rknn

app = Flask(__name__)


def pz(image1):
    image1 = cv2.flip(image1, 0)
    tx =-10# 左右平移距离，正值向右移动，负值向左移动
    ty = -45  # 上下平移距离，正值向下移动，负值向上移动
    # 定义水平和垂直拉伸因子
    scale_x = 1.2  # 水平拉伸因子，大于1表示拉伸，小于1表示压缩
    scale_y = 1  # 垂直拉伸因子，大于1表示拉伸，小于1表示压缩
    # 构建仿射变换矩阵
    affine_matrix = np.array([[scale_x, 0, tx],
                              [0, scale_y, ty]], dtype=np.float32)
    # 应用仿射变换
    return cv2.warpAffine(image1, affine_matrix, (image1.shape[1], image1.shape[0]))

def remix(img1, img2, alpha=0.3):
    result = cv2.addWeighted(img1, 1 - alpha, img2, alpha, 0)
    cv2.imwrite('./vis.jpg', img1)
    cv2.imwrite('./ir.jpg', img2)
    return result

def generate_frames(cam1,cam2):
    
    while True:
        ret1,img1=cam1.get_frame()
        ret2,img2=cam2.get_frame()
        if ret1 and ret2:
            img1 = cv2.resize(img1, (640, 480))
            img2 = cv2.resize(img2, (640, 480))
            img1=pz(img1)
            # result = remix(img1,img2,0.8)
            result = rm.remix(img1,img2)
            ret, buffer = cv2.imencode('.jpg', result)
            frame = buffer.tobytes()
            yield (b'--frame\r\n'
                    b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')

@app.route('/')
def video_feed():
    return Response(generate_frames(cam1,cam2),
                    mimetype='multipart/x-mixed-replace; boundary=frame')

if __name__ == '__main__':
    cam1=Camera(camera_id=0)
    cam2=Camera(camera_id=1)
    rm=infer_rknn.remix_model()
    cam1.start()
    cam2.start()
    app.run(host='0.0.0.0', port=5000)